#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：yolo_pipeline.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/1/7 上午9:22
"""

import os
from yolo_data_builder import YOLOv5DataBuilder
from yolo_train import YOLOv5Trainer
from yolo_test import YOLOv5Tester


class YOLOv5Pipeline:
	def __init__(self, input_dir: str, output_dir: str, train_ratio: float = 0.8):
		"""
		初始化YOLOv5流程。

		:param input_dir: 包含images和labels文件夹的输入目录路径。
		:param output_dir: 输出YOLO格式数据的目录路径。
		:param train_ratio: 训练集和测试集的比例，默认为0.8。
		"""
		self.input_dir = os.path.abspath(input_dir)
		self.output_dir = os.path.abspath(output_dir)
		self.train_ratio = train_ratio

	def run(self):
		"""
		运行YOLOv5流程。
		"""
		# 数据准备
		data_builder = YOLOv5DataBuilder(self.input_dir, self.output_dir, self.train_ratio)
		data_builder.prepare_dataset()

		# 获取类别列表
		classes = data_builder.classes

		# 训练模型
		train_images_dir = os.path.join(self.output_dir, 'train')
		trainer = YOLOv5Trainer(train_images_dir, self.output_dir, classes)
		trainer.train_model()

		# 测试模型
		test_images_dir = os.path.join(self.output_dir, 'test')
		tester = YOLOv5Tester(test_images_dir, self.output_dir, classes)
		tester.test_model()


# 使用示例
if __name__ == "__main__":
	input_dir = 'datasets/HELMET/source_data'
	output_dir = 'datasets/HELMET/yolo_output_dir'
	pipeline = YOLOv5Pipeline(input_dir, output_dir)
	pipeline.run()
